#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""  
@Project : PythonProject
@File : login_widget.py.py
@Author : donc
@Time : 2025/5/31 2:13  
@脚本说明 : MySQL数据库登录组件
"""
import streamlit as st
import datetime
from database import get_user, verify_password, update_last_login
from utils import set_background  # 确保有这个工具函数
from extra_streamlit_components import CookieManager  # 添加安全 Cookie 管理


class LoginWidget:
    """
    基于MySQL数据库的登录组件
    """

    COOKIE_NAME = "__mysql_login_username__"

    def __init__(self, width=200, height=280, lottie_url="https://assets8.lottiefiles.com/packages/lf20_ktwnwv5m.json"):
        self.width = width
        self.height = height
        self.lottie_url = lottie_url
        self.cookie_manager = CookieManager(key="auth_cookies")  # 使用安全的 Cookie 管理器

        # 初始化会话状态
        if 'LOGOUT_BUTTON_HIT' not in st.session_state:
            st.session_state.LOGOUT_BUTTON_HIT = False
        if 'logged_in' not in st.session_state:
            st.session_state.logged_in = False
        if 'user' not in st.session_state:
            st.session_state.user = None
        if 'LOGGED_IN_USERNAME' not in st.session_state:
            st.session_state.LOGGED_IN_USERNAME = None
        if 'LOGGED_IN_EMAIL' not in st.session_state:
            st.session_state.LOGGED_IN_EMAIL = None
        if 'LOGGED_IN_ROLE' not in st.session_state:
            st.session_state.LOGGED_IN_ROLE = None

    def _get_cookie(self) -> str:
        """从查询参数获取Cookie值（使用新的 st.query_params API）"""
        return self.cookie_manager.get(self.COOKIE_NAME)

    def _set_cookie(self, value: str, expires_days: int = 7):
        """设置安全的 Cookie"""
        self.cookie_manager.set(
            self.COOKIE_NAME,
            value,
            expires_at=datetime.datetime.now() + datetime.timedelta(days=expires_days)
        )

    def _delete_cookie(self):
        """删除 Cookie"""
        if self.COOKIE_NAME in self.cookie_manager.cookies:
            self.cookie_manager.delete(self.COOKIE_NAME)

    def get_username(self):
        """从Cookie获取用户名"""
        if st.session_state.LOGOUT_BUTTON_HIT == False:
            username = self._get_cookie()
            if username and username != 'dummy':
                return username
        return None

    def login_widget(self):
        """登录表单组件"""
        if st.session_state.logged_in:
            return

        # 设置登录页面背景
        set_background()

        # 双列布局
        col1, col2, col3 = st.columns([3.5, 3, 3.5])

        with col2:
            # 登录卡片容器
            st.markdown("""
            <style>
            .login-card {
                background-color: rgba(255, 255, 255, 0.9);
                border-radius: 15px;
                padding: 2rem;
                box-shadow: 0 10px 30px rgba(0, 0, 0, 0.2);
                margin-top: 5rem;
            }
            .login-title {
                text-align: center;
                color: #2c3e50;
                margin-bottom: 1.5rem;
            }
            </style>
            """, unsafe_allow_html=True)

            with st.container():
                # st.markdown('<div class="login-card">', unsafe_allow_html=True)

                st.markdown('<h2 class="login-title">智慧安全漏扫平台</h2>', unsafe_allow_html=True)


                with st.form("Login Form", clear_on_submit=True):
                    username = st.text_input("用户名", placeholder='输入用户名', key="login_username")
                    password = st.text_input("密码", placeholder='输入密码', type='password', key="login_password")

                    st.markdown("###")
                    login_submit_button = st.form_submit_button("登录", use_container_width=True)

                    if login_submit_button:
                        if not username or not password:
                            st.error("请输入用户名和密码")
                            return

                        user = get_user(username)
                        if user and verify_password(password, user['password']):
                            st.session_state.logged_in = True
                            st.session_state.LOGOUT_BUTTON_HIT = False

                            # 更新会话状态
                            st.session_state.user = user
                            st.session_state.LOGGED_IN_USERNAME = user['username']
                            st.session_state.LOGGED_IN_EMAIL = user.get('email', '')
                            st.session_state.LOGGED_IN_ROLE = user.get('role', 'user')

                            # 不再设置Cookie
                            self._delete_cookie()

                            # 更新最后登录时间
                            update_last_login(username)
                            st.rerun()
                        else:
                            st.error("用户名或密码错误")

                st.markdown('</div>', unsafe_allow_html=True)

    def logout_widget(self):
        """登出组件"""
        if st.session_state.logged_in:
            unique_key = "logout_btn"
            if st.sidebar.button("登出", key=unique_key, use_container_width=True):
                # 清除所有会话状态
                keys_to_clear = [
                    "LOGOUT_BUTTON_HIT", "logged_in", "user",
                    "LOGGED_IN_USERNAME", "LOGGED_IN_EMAIL", "LOGGED_IN_ROLE",
                    "mysql_scan_result", "redis_scan_result", "host_scan_result",
                    "port_scan_result", "ssh_scan_result", "show_admin_panel"
                ]
                for key in keys_to_clear:
                    if key in st.session_state:
                        del st.session_state[key]

                self._delete_cookie()
                st.rerun()

    def build_login_ui(self):
        """构建登录UI并返回登录状态"""
        self.login_widget()
        return st.session_state.logged_in